#!/usr/bin/env python3

import argparse
import pymongo
from pprint import pprint

parser = argparse.ArgumentParser(description='Calculate observation neighborhood and save it to database.')
parser.add_argument('--mongo-uri', help='Observation database URI', required=True)
parser.add_argument('--obs-type', help='Observation type', choices=('synop', 'raob', 'amdar', 'profiler'), required=True)
parser.add_argument('--radius', help='Observation site group radius (km)', required=True, type=float)
args = parser.parse_args()

col_names = {
	'synop'   : 'synop_stations',
	'raob'    : 'raob_stations',
	'amdar'   : 'amdar_records',
	'profiler': 'profiler_stations'
}

client = pymongo.MongoClient(f'mongodb://{args.mongo_uri}')

db = client.metobs
col_site = db[col_names[args.obs_type]]

sid_list = col_site.distinct('sid')

col_ngb = db[f'{args.obs_type}_ngb']

col_ngb.create_index([( 'sid', 'text' )], unique=True)

ngb = []
for sid in sid_list:
	print(f'==> {sid} ...')
	loc = col_site.find_one({ 'sid': sid })['location']
	ngb.append({
		'sid': sid,
		'ngb': list(col_site.aggregate([
			{
				'$geoNear': {
					'near': loc,
					'maxDistance': args.radius * 1000, # m
					'distanceField': 'd',
					'spherical': True
				}
			}
		]))
	})
	if len(ngb) >= 100:
		print(f'Insert {len(ngb)} neighborhood data ...')
		try:
			col_ngb.insert_many(ngb)
		except pymongo.errors.BulkWriteError as e:
			if e.details['writeErrors'][0]['code'] == 11000: # Duplicate key, just skip it.
				pass
			else:
				print(f'[Error]: Failed to write database! See error: {e.details["writeErrors"][0]["errmsg"]}')
				exit(1)
		ngb = []

if len(ngb) > 0:
	print(f'Insert {len(ngb)} neighborhood data ...')
	col_ngb.insert_many(ngb)
